import torch
from torch import nn
from fasttext import FTModel
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv, GCNConv
from torch_geometric.nn import GATConv, Linear, to_hetero


class Classifier(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(Classifier, self).__init__()
        self.project = nn.Sequential(
            # nn.BatchNorm1d(input_size),
            nn.Linear(input_size, hidden_size, bias=True),
            # nn.LayerNorm(hidden_size),
            # nn.Tanh(),
            nn.LeakyReLU(0.2),
            # nn.LayerNorm(),
            # nn.Linear(hidden_size, hidden_size,bias=True),
            # nn.BatchNorm1d(hidden_size)
        )

        self.project2 = nn.Sequential(
            # nn.BatchNorm1d(input_size),
            nn.Linear(input_size, hidden_size, bias=True),
            # nn.LayerNorm(hidden_size),
            # nn.Tanh(),
            nn.LeakyReLU(0.2),
            # nn.LayerNorm(),
            # nn.Linear(hidden_size, hidden_size,bias=True),
            # nn.BatchNorm1d(hidden_size)
        )

    def forward(self, x, edge_label_index):
        x_alarm = x["alarm"]
        x_btree = x["btree"]
        x_btree = self.project(x_btree)
        x_alarm = self.project2(x_alarm)
        edge_feat_A = x_alarm[edge_label_index[0]]
        edge_feat_B = x_btree[edge_label_index[1]]
        out = (edge_feat_A * edge_feat_B).sum(dim=-1)
        # out = torch.sigmoid(out)
        return out


# class GAT_Model(nn.Module):
#     def __init__(self, hidden_channels):
#         super(GAT_Model, self).__init__()
#         self.conv1 = GATConv(hidden_channels, hidden_channels, heads=4)
#         self.conv2 = GATConv(hidden_channels * 4, hidden_channels)
#
#     def forward(self, x, edge_index):
#         x = F.relu(self.conv1(x, edge_index))
#         x = self.conv2(x, edge_index)
#         return x


class GAT_Model(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        # self.conv1 = GATConv((-1, -1), hidden_channels , add_self_loops=False, heads=4)
        # self.lin1 = Linear(-1, hidden_channels)
        # self.conv2 = GATConv((-1, -1), out_channels , add_self_loops=False, heads=4)
        self.conv1 = GATConv((-1, -1), hidden_channels // 4, add_self_loops=False, heads=4)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels // 4, add_self_loops=False, heads=4)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x


class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, 1)

    def forward(self, z_dict, edge_label_index):
        row, col = edge_label_index
        z = torch.cat([z_dict['alarm'][row], z_dict['btree'][col]], dim=-1)

        z = self.lin1(z).relu()
        z = self.lin2(z)
        return torch.sigmoid(z.view(-1))


class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x


class Model(nn.Module):
    def __init__(self, hidden_channels, data):
        super().__init__()
        self.encoder = GAT_Model(hidden_channels, hidden_channels)
        self.encoder = to_hetero(self.encoder, metadata=data.metadata(), aggr='sum')

        self.decoder = EdgeDecoder(hidden_channels)

    def forward(self, batch_data):
        x_dict = {
            "alarm": batch_data["alarm"].x,
            "btree": batch_data["btree"].x,
        }
        new_x = self.encoder(x_dict, batch_data.edge_index_dict)
        out = self.decoder(new_x, batch_data[("alarm", "to", "btree")].edge_label_index)
        return out
